package org.sippr.utils;

import com.fasterxml.jackson.core.JsonProcessingException;
import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.fasterxml.jackson.databind.node.ArrayNode;
import org.sippr.pojo.SdPara;
import org.sippr.pojo.TaskQueue;
import org.sippr.service.impl.HttpSdModels;

import java.awt.image.BufferedImage;
import java.nio.file.Path;
import java.nio.file.Paths;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

public class SdRequestUtils {


    public static Boolean SdCreateImgBySkech(String baseUrl, TaskQueue task, String pngPath, String baseDirectory) {
        String url = baseUrl + "/sdapi/v1/txt2img";
        Path dir = Paths.get(baseDirectory, "\\skech");
        String filename = task.getCreationImgSketch().getName();
        Path skechPath = dir.resolve(filename);
        String imgBase64 = ImgUtil.imgToBase64(skechPath.toString());
        SdPara sdPara = JsonUtil.parseJson(task.getCreationStyle().getParas(), SdPara.class);


        List<Map<String, Object>> args = new ArrayList<>();

        Map<String, Object> arg = new HashMap<>();
        arg.put("enabled", "True");
        arg.put("module", "lineart");
        arg.put("model", "control_v11p_sd15_lineart [43d4be0d]");
        arg.put("weight", 1.0);
        arg.put("control_mode", sdPara.getControlMode());
        arg.put("resize_mode", "Crop and Resize");
        arg.put("image", imgBase64);
        arg.put("resize_mode", 1);
        arg.put("lowvram", "False");
        arg.put("processor_res", 512);
        arg.put("threshold_a", 64);
        arg.put("threshold_b", 64);
        arg.put("guidance_start", 0.0);
        arg.put("guidance_end", 1.0);
        arg.put("control_mode", 0);
        arg.put("pixel_perfect", "False");
        args.add(arg);
        Map<String, Object> controlnet = new HashMap<>();
        controlnet.put("args", args);


        Map<String, Object> alwayson_scripts = new HashMap<>();
        alwayson_scripts.put("controlnet", controlnet);


        Map<String, Object> body = new HashMap<>();
        body.put("prompt", task.getCreationWordCue().getWords() + "," + task.getSdCheckPoints().getWordCue());
        body.put("negative_prompt", task.getCreationWordNegative().getWords() + "," + task.getSdCheckPoints().getWordNega());
        body.put("sampler_index", "Euler a");
        body.put("batch_size", 1);
        body.put("steps", sdPara.getSamplingSteps());
        body.put("cfg_scale", sdPara.getCfgScale());
        body.put("height", sdPara.getHeight());
        body.put("width", sdPara.getWidth());
        body.put("alwayson_scripts", alwayson_scripts);

        String json = HttpUtil.PostByJson(url, body);

        ObjectMapper objectMapper = new ObjectMapper();
        try {
            JsonNode rootNode = objectMapper.readTree(json);
            ArrayNode hobbiesArrayNode = (ArrayNode) rootNode.get("images");
            if (hobbiesArrayNode != null) {
                for (JsonNode hobbyNode : hobbiesArrayNode) {
                    String hobby = hobbyNode.asText();
                    BufferedImage img = ImgUtil.processBase64Image(hobby);
                    ImgUtil.saveImage(img, pngPath);
                    return true;

                }
            }
        } catch (Exception e) {
            throw new RuntimeException(e);
        }
        return false;
    }


    //从sd服务获取已有的lorra模型列表
    public static String GetLorras(String baseUrl) {
        String url = baseUrl + "/sdapi/v1/loras";
        return HttpUtil.Get(url, null);

    }

    //从sd服务获取已有的cpt模型列表
    public static HttpSdModels[] GetCpts(String baseUrl) throws JsonProcessingException {
        String url = baseUrl + "/sdapi/v1/sd-models";
        String json = HttpUtil.Get(url, null);
        ObjectMapper objectMapper = new ObjectMapper();
        return objectMapper.readValue(json, HttpSdModels[].class);
    }

    //从sd服务获取已选用的sd大模型
    public static String GetSdSelectedCpt(String baseUrl) throws JsonProcessingException {
        //sd_model_checkpoint:"architecturerealmix_v1repair.safetensors [f7c2a4bb41]"
        String url = baseUrl + "/sdapi/v1/options";
        String json = HttpUtil.Get(url, null);
        ObjectMapper objectMapper = new ObjectMapper();
        JsonNode rootNode = objectMapper.readTree(json);
        return rootNode.get("sd_model_checkpoint").asText();
    }

    public static String GetSdProgress(String baseUrl) {
        String url = baseUrl + "/sdapi/v1/progress";
        return HttpUtil.Get(url, null);
    }

    //判断SD是否在线，即是否能够访问
    public static Boolean IsOnLine(String baseUrl) {
        String url = baseUrl + "/sdapi/v1/progress";
        return HttpUtil.IsOnLine(url);
    }

    //判断SD是否繁忙，即是否工作中
    public static Boolean IsBusy(String baseUrl) {
        try {
            String json = GetSdProgress(baseUrl);
            ObjectMapper objectMapper = new ObjectMapper();
            JsonNode rootNode = objectMapper.readTree(json);
            Double p = rootNode.get("progress").asDouble();
            if (p == 0.0)
                return false;
            else
                return true;
        } catch (Exception e) {
            System.out.println(e);
            return true;
        }
    }

    public static String SdQueueStatus(String baseUrl) {
        String url = baseUrl + "/queue/status";
        return HttpUtil.Get(url, null);
    }

    public static Void SdChangeCpt(String baseUrl, String cpAlias) throws JsonProcessingException {
        //sd_model_checkpoint:"architecturerealmix_v1repair.safetensors [f7c2a4bb41]"
        String url = baseUrl + "/sdapi/v1/options";
        String selectedCp = GetSdSelectedCpt(baseUrl);
        HttpSdModels[] httpSdModels = GetCpts(baseUrl);
        for (HttpSdModels httpSdModel : httpSdModels) {
            if (httpSdModel.getModel_name().contains(cpAlias)) {
                if (httpSdModel.getTitle() != selectedCp) {
                    Map<String, Object> body = new HashMap<>();
                    body.put("sd_model_checkpoint", httpSdModel.getTitle());
                    body.put("sd_vae", "Automatic");
                    HttpUtil.PostByJson(url, body);

                }
                break;
            }
        }

        return null;
    }

    public static Boolean SdCreateImgByTxt(String baseUrl, TaskQueue task, String pngPath) {
        String url = baseUrl + "/sdapi/v1/txt2img";
        SdPara sdPara = JsonUtil.parseJson(task.getCreationStyle().getParas(), SdPara.class);

        Map<String, Object> body = new HashMap<>();
        body.put("batch_size", 1);
        body.put("cfg_scale", sdPara.getCfgScale());
        body.put("denoising_strength", 0);
        body.put("enable_hr", "False");
        body.put("eta", 0);
//        body.put("firstphase_height", 0);
//        body.put("firstphase_width", 0);
        body.put("height", sdPara.getHeight());
        body.put("n_iter", 1);
        body.put("negative_prompt", task.getCreationWordNegative().getWords() + "," + task.getSdCheckPoints().getWordNega());
        body.put("prompt", task.getCreationWordCue().getWords() + "," + task.getSdCheckPoints().getWordCue());
        body.put("restore_faces", "False");
        body.put("s_churn", 0);
        body.put("s_noise", 1);
        body.put("s_tmax", 0);
        body.put("s_tmin", 0);
//        body.put("sampler_index", "Euler a");
        body.put("sampler_index", "DPM++ 2M Karras");
        body.put("seed", -1);
        body.put("seed_resize_from_h", -1);
        body.put("seed_resize_from_w", -1);
        body.put("steps", sdPara.getSamplingSteps());
//        body.put(   "styles", []);
        body.put("subseed", -1);
        body.put("subseed_strength", 0);
        body.put("tiling", "False");
        body.put("width", sdPara.getWidth());

        String json = HttpUtil.PostByJson(url, body);

        ObjectMapper objectMapper = new ObjectMapper();
        System.out.println("Response Body : " + body.toString());
        try {
            JsonNode rootNode = objectMapper.readTree(json);
            ArrayNode hobbiesArrayNode = (ArrayNode) rootNode.get("images");
            if (hobbiesArrayNode != null) {
                for (JsonNode hobbyNode : hobbiesArrayNode) {
                    String hobby = hobbyNode.asText();
                    BufferedImage img = ImgUtil.processBase64Image(hobby);
                    ImgUtil.saveImage(img, pngPath);
                    return true;

                }
            }
        } catch (Exception e) {
            throw new RuntimeException(e);
        }
        return false;
    }


    //获取controlnet预览图
    public static Boolean SdControlNetDetect(String baseUrl, String srcPngPath, String destPngPath) {
        String url = baseUrl + "/controlnet/detect";
        //获取原图片BASE64编码
        String imgBase64 = ImgUtil.imgToBase64(srcPngPath);
        List<String> img64List = new ArrayList<>();
        img64List.add(imgBase64);
        Map<String, Object> body = new HashMap<>();
        body.put("controlnet_module", "lineart_standard (from white bg & black line)");
        body.put("controlnet_input_images", img64List);

        String json = HttpUtil.PostByJson(url, body);

        ObjectMapper objectMapper = new ObjectMapper();
        try {
            JsonNode rootNode = objectMapper.readTree(json);
            ArrayNode hobbiesArrayNode = (ArrayNode) rootNode.get("images");
            if (hobbiesArrayNode != null) {
                for (JsonNode hobbyNode : hobbiesArrayNode) {
                    String hobby = hobbyNode.asText();
                    BufferedImage img = ImgUtil.processBase64Image(hobby);
                    ImgUtil.saveImage(img, destPngPath);
                    return true;
                }
            }
        } catch (Exception e) {
            throw new RuntimeException(e);
        }
        return false;
    }

}
